import cv2
import numpy as np


def get_lines(edge_img):
    """
    获取edge_img中的所有线段
    :param edge_img: 标记边缘的灰度图
    """

    def calculate_slope(line):
        """
        计算线段line的斜率
        :param line: np.array([[x_1, y_1, x_2, y_2]])
        :return:
        """
        x_1, y_1, x_2, y_2 = line[0]
        return (y_2 - y_1) / (x_2 - x_1)

    def reject_abnormal_lines(lines, threshold=0.2):
        """
        剔除斜率不一致的线段
        :param lines: 线段集合, [np.array([[x_1, y_1, x_2, y_2]]),np.array([[x_1, y_1, x_2, y_2]]),...,np.array([[x_1, y_1, x_2, y_2]])]
        """
        slopes = [calculate_slope(line) for line in lines]
        while len(lines) > 0:
            mean = np.mean(slopes)
            diff = [abs(s - mean) for s in slopes]
            idx = np.argmax(diff)
            if diff[idx] > threshold:
                slopes.pop(idx)
                lines.pop(idx)
            else:
                break
        return lines

    def least_squares_fit(lines):
        """
        将lines中的线段拟合成一条线段
        :param lines: 线段集合, [np.array([[x_1, y_1, x_2, y_2]]),np.array([[x_1, y_1, x_2, y_2]]),...,np.array([[x_1, y_1, x_2, y_2]])]
        :return: 线段上的两点,np.array([[xmin, ymin], [xmax, ymax]])
        """
        x_coords = np.ravel([[line[0][0], line[0][2]] for line in lines])
        y_coords = np.ravel([[line[0][1], line[0][3]] for line in lines])
        poly = np.polyfit(x_coords, y_coords, deg=1)
        point_min = (np.min(x_coords), np.polyval(poly, np.min(x_coords)))
        point_max = (np.max(x_coords), np.polyval(poly, np.max(x_coords)))
        return np.array([point_min, point_max], dtype=np.int)

    # 获取所有线段
    lines = cv2.HoughLinesP(edge_img, 1, np.pi / 180, 15, minLineLength=40,
                            maxLineGap=20)
    # 按照斜率分成车道线
    left_lines = [line for line in lines if calculate_slope(line) > 0]
    right_lines = [line for line in lines if calculate_slope(line) < 0]
    # 剔除离群线段
    left_lines = reject_abnormal_lines(left_lines)
    right_lines = reject_abnormal_lines(right_lines)

    return least_squares_fit(left_lines), least_squares_fit(right_lines)

# 读取图像
image = cv2.imread('images/2.jpg')

# 将图像从BGR颜色空间转换为HSV颜色空间
hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

# 设置颜色上下阈值（这里以蓝色为例）
lower_blue = np.array([0,100,100])
upper_blue = np.array([10, 255, 255])

# 应用阈值化，获取颜色区域的掩膜
color_mask = cv2.inRange(hsv_image, lower_blue, upper_blue)

# 对原始图像和掩膜进行位运算，提取颜色区域
color_region = cv2.bitwise_and(image, image, mask=color_mask)

# 边缘检测
edges = cv2.Canny(color_region, 50, 150, apertureSize=3)

# 使用霍夫变换找到线段
lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)

print(lines)

# 计算每条线段的斜率
slopes = []
for rho, theta in lines[:, 0]:
    # 将弧度转换为度
    angle = np.degrees(theta)
    # 如果角度大于90度，调整斜率的方向
    if angle > 90:
        angle -= 180
    # 计算直线的斜率
    if angle != 0:
        slope = np.tan(np.radians(angle))
        slopes.append(slope)

# 输出斜率列表
print(slopes)

# 显示提取的颜色区域
cv2.imshow("Color Region", color_region)
cv2.waitKey(0)
cv2.destroyAllWindows()


